package com.zw.cc.redis;

import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.Nullable;
import org.springframework.util.CollectionUtils;
import redis.clients.jedis.*;
import redis.clients.jedis.exceptions.JedisConnectionException;
import redis.clients.jedis.exceptions.JedisException;
import redis.clients.util.Hashing;
import redis.clients.util.Pool;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.regex.Pattern;

public class ShardedJedisSentinelPool extends Pool<ShardedJedis> {

    public static final int MAX_RETRY_SENTINEL = 10;

    protected final Logger log = LoggerFactory.getLogger(getClass().getName());

    protected GenericObjectPoolConfig poolConfig;

    protected int connectionTimeout = Protocol.DEFAULT_TIMEOUT;

    private int sentinelRetry = 0;

    protected String password;

    protected int database = Protocol.DEFAULT_DATABASE;

    protected Set<MasterListener> masterListeners = new HashSet<>();

    private volatile List<HostAndPort> currentHostMasters;

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, final GenericObjectPoolConfig poolConfig) {
        this(masterNames, sentinels, poolConfig, Protocol.DEFAULT_TIMEOUT, null, Protocol.DEFAULT_DATABASE);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels) {
        this(masterNames, sentinels, new GenericObjectPoolConfig(), Protocol.DEFAULT_TIMEOUT, null, Protocol.DEFAULT_DATABASE);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, String password) {
        this(masterNames, sentinels, new GenericObjectPoolConfig(), Protocol.DEFAULT_TIMEOUT, password);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, final GenericObjectPoolConfig poolConfig, int timeout,
                                    final String password) {
        this(masterNames, sentinels, poolConfig, timeout, password, Protocol.DEFAULT_DATABASE);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, final GenericObjectPoolConfig poolConfig, final int timeout) {
        this(masterNames, sentinels, poolConfig, timeout, null, Protocol.DEFAULT_DATABASE);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, final GenericObjectPoolConfig poolConfig, final String password) {
        this(masterNames, sentinels, poolConfig, Protocol.DEFAULT_TIMEOUT, password);
    }

    public ShardedJedisSentinelPool(List<String> masterNames, Set<String> sentinels, final GenericObjectPoolConfig poolConfig, int timeout,
                                    final String password, final int database) {
        this.poolConfig = poolConfig;
        this.connectionTimeout = timeout;
        this.password = password;
        this.database = database;

        List<HostAndPort> masterList = initSentinels(sentinels, masterNames);
        initPool(masterList);
    }

    @Override
    public void destroy() {
        for (MasterListener m : masterListeners) {
            m.shutdown();
        }

        super.destroy();
    }

    public List<HostAndPort> getCurrentHostMaster() {
        return currentHostMasters;
    }

    private void initPool(List<HostAndPort> masters) {
        if (!equals(currentHostMasters, masters)) {
            StringBuilder sb = new StringBuilder();
            for (HostAndPort master : masters) {
                sb.append(master.toString());
                sb.append(" ");
            }
            log.info("Created ShardedJedisPool to master at [" + sb.toString() + "]");
            initPool(poolConfig, new ShardedJedisFactory(makeShardInfoList(masters), Hashing.MURMUR_HASH, null));
            currentHostMasters = masters;
        }
    }

    private boolean equals(List<HostAndPort> currentShardMasters, List<HostAndPort> shardMasters) {
        if (currentShardMasters != null && shardMasters != null && currentShardMasters.size() == shardMasters.size()) {
            for (int i = 0; i < currentShardMasters.size(); i++) {
                if (!currentShardMasters.get(i).equals(shardMasters.get(i)))
                    return false;
            }
            return true;
        }
        return false;
    }

    private List<JedisShardInfo> makeShardInfoList(List<HostAndPort> masters) {
        List<JedisShardInfo> shardMasters = new ArrayList<>();
        for (HostAndPort master : masters) {
            JedisShardInfo jedisShardInfo = new JedisShardInfo(master.getHost(), master.getPort(), connectionTimeout);
            jedisShardInfo.setPassword(password);

            shardMasters.add(jedisShardInfo);
        }
        return shardMasters;
    }

    /**
     * 初始化 哨兵
     *
     * @param sentinels   哨兵IP地址
     * @param masterNames master
     * @return
     */
    private List<HostAndPort> initSentinels(@Nullable Set<String> sentinels, @Nullable final List<String> masterNames) {
        if (CollectionUtils.isEmpty(sentinels)) {
            log.error("No sentinels configed, Abort. ");
            throw new JedisException("No sentinels configed, Abort. ");
        }
        if (CollectionUtils.isEmpty(masterNames)) {
            log.error("No master config, Abort. ");
            throw new JedisException("Cannot connect all sentinels, No master config, Abort.");
        }
        Map<String, HostAndPort> masterMap = new HashMap<>();
        List<HostAndPort> shardMasters = new ArrayList<>();

        log.info("Trying to find all master from available Sentinels...");

        for (String masterName : masterNames) {
            HostAndPort master = null;
            boolean fetched = false;
            while (!fetched && sentinelRetry < MAX_RETRY_SENTINEL) {
                for (String sentinel : sentinels) {
                    final HostAndPort hap = toHostAndPort(Arrays.asList(sentinel.split(":")));

                    log.info("Connecting to Sentinel " + hap);

                    try (Jedis jedis = new Jedis(hap.getHost(), hap.getPort())) {
                        master = masterMap.get(masterName);
                        if (master == null) {
                            List<String> hostAndPort = jedis.sentinelGetMasterAddrByName(masterName);
                            if (hostAndPort != null && !hostAndPort.isEmpty()) {
                                master = toHostAndPort(hostAndPort);
                                log.info("Found Redis master at " + master);
                                shardMasters.add(master);
                                masterMap.put(masterName, master);
                                fetched = true;
                            }
                        }
                    } catch (JedisConnectionException e) {
                        log.warn("Cannot connect to sentinel running @ " + hap + ". Trying next one.");
                    }
                    if (fetched)
                        break;
                }

                if (null == master) {
                    try {
                        log.error("All sentinels down, cannot determine where is " + masterName + " master is running... sleeping 1000ms, Will try again.");
                        Thread.sleep(1000);
                    } catch (InterruptedException e) {
                        log.error("Thread.sleep(1000) InterruptedException",e);
                        Thread.currentThread().interrupt();
                    }
                    fetched = false;
                    sentinelRetry++;
                }
            }

            if (!fetched && sentinelRetry >= MAX_RETRY_SENTINEL) {
                log.error("All sentinels down and try " + MAX_RETRY_SENTINEL + " times, Abort.");
                throw new JedisConnectionException("Cannot connect all sentinels, Abort.");
            }
        }

        if (masterNames.size() == shardMasters.size()) {

            log.info("Starting Sentinel listeners...");
            for (String sentinel : sentinels) {
                final HostAndPort hap = toHostAndPort(Arrays.asList(sentinel.split(":")));
                MasterListener masterListener = new MasterListener(masterNames, hap.getHost(), hap.getPort());
                masterListener.setDaemon(true);
                masterListeners.add(masterListener);
                masterListener.start();
            }
        }

        return shardMasters;
    }

    private HostAndPort toHostAndPort(List<String> getMasterAddrByNameResult) {
        String host = getMasterAddrByNameResult.get(0);
        int port = Integer.parseInt(getMasterAddrByNameResult.get(1));

        return new HostAndPort(host, port);
    }

    @Override
    public ShardedJedis getResource() {
        while (true) {
            final ShardedJedis shardedJedis = super.getResource();
            shardedJedis.setDataSource(this);

            final List<HostAndPort> curHostMasters = currentHostMasters;
            final List<HostAndPort> connections = new ArrayList<HostAndPort>();
            shardedJedis.getAllShards().forEach(jedis ->
                    connections.add(new HostAndPort(jedis.getClient().getHost(), jedis.getClient().getPort())));
            if (curHostMasters.size() == connections.size()) {
                for (HostAndPort host : curHostMasters) {
                    if (!connections.contains(host))
                        returnBrokenResource(shardedJedis);
                }
                return shardedJedis;
            } else {
                returnBrokenResource(shardedJedis);
            }
        }
    }

    @Override
    public void returnBrokenResource(final ShardedJedis resource) {
        if (resource != null) {
            returnBrokenResourceObject(resource);
        }
    }

    @SuppressWarnings("deprecation")
    @Override
    public void returnResource(final ShardedJedis resource) {
        if (resource != null) {
            resource.resetState();
            internalPool.returnObject(resource);
        }
    }

    public void returnResource(final ShardedJedis resource, final boolean broken) {
        if (broken) {
            returnBrokenResource(resource);
        } else {
            returnResource(resource);
        }
    }

    // =================================================================
    protected class MasterListener extends Thread {

        protected List<String> masterNames;

        protected String host;

        protected int port;

        protected long subscribeRetryWaitTimeMillis = 5000;

        protected volatile Jedis j;

        protected AtomicBoolean running = new AtomicBoolean(false);

        protected MasterListener() {
        }

        public MasterListener(List<String> masterNames, String host, int port) {
            this.masterNames = masterNames;
            this.host = host;
            this.port = port;
        }

        public MasterListener(List<String> masterNames, String host, int port, long subscribeRetryWaitTimeMillis) {
            this(masterNames, host, port);
            this.subscribeRetryWaitTimeMillis = subscribeRetryWaitTimeMillis;
        }
        @Override
        public void run() {

            running.set(true);

            while (running.get()) {

                j = new Jedis(host, port);

                try {
                    // double check that it is not being shutdown
                    if (!running.get()) {
                        break;
                    }

                    j.subscribe(new JedisPubSub() {

                        @Override
                        public void onMessage(String channel, String message) {
                            log.info("Sentinel " + host + ":" + port + " published: " + message + ".");

                            String[] switchMasterMsg = message.split(" ");

                            if (switchMasterMsg.length > 3) {

                                int index = masterNames.indexOf(switchMasterMsg[0]);
                                if (index >= 0) {
                                    HostAndPort newHostMaster = toHostAndPort(Arrays.asList(switchMasterMsg[3], switchMasterMsg[4]));
                                    List<HostAndPort> newHostMasters = new ArrayList<HostAndPort>();
                                    for (int i = 0; i < masterNames.size(); i++) {
                                        newHostMasters.add(null);
                                    }
                                    Collections.copy(newHostMasters, currentHostMasters);
                                    newHostMasters.set(index, newHostMaster);

                                    initPool(newHostMasters);
                                } else {
                                    StringBuilder sb = new StringBuilder();
                                    for (String masterName : masterNames) {
                                        sb.append(masterName);
                                        sb.append(",");
                                    }
                                    log.info("Ignoring message on +switch-master for master name " + switchMasterMsg[0] + ", our monitor master name are ["
                                            + sb + "]");
                                }

                            } else {
                                log.error("Invalid message received on Sentinel " + host + ":" + port + " on channel +switch-master: " + message);
                            }
                        }
                    }, "+switch-master");

                } catch (JedisConnectionException e) {

                    if (running.get()) {
                        log.info("Lost connection to Sentinel at " + host + ":" + port + ". Sleeping 5000ms and retrying.", e);
                        try {
                            Thread.sleep(subscribeRetryWaitTimeMillis);
                        } catch (InterruptedException e1) {
                            log.info("Sleep interrupted: ", e1);
                            Thread.currentThread().interrupt();
                        }
                    } else {
                        log.info("Unsubscribing from Sentinel at " + host + ":" + port);
                    }
                } finally {
                    j.close();
                }
            }
        }

        public void shutdown() {
            try {
                log.info("Shutting down listener on " + host + ":" + port);
                running.set(false);
                // This isn't good, the Jedis object is not thread safe
                if (j != null) {
                    j.disconnect();
                }
            } catch (Exception e) {
                log.info("Caught exception while shutting down: ", e);
            }
        }
    }

}

/**
 * PoolableObjectFactory custom impl.
 */
class ShardedJedisFactory implements PooledObjectFactory<ShardedJedis> {

    public static final Logger log = LoggerFactory.getLogger(ShardedJedisFactory.class.getName());

    private List<JedisShardInfo> shards;

    private Hashing algo;

    private Pattern keyTagPattern;

    public ShardedJedisFactory(List<JedisShardInfo> shards, Hashing algo, Pattern keyTagPattern) {
        this.shards = shards;
        this.algo = algo;
        this.keyTagPattern = keyTagPattern;
    }

    public PooledObject<ShardedJedis> makeObject() throws Exception {
        ShardedJedis jedis = new ShardedJedis(shards, algo, keyTagPattern);
        return new DefaultPooledObject<ShardedJedis>(jedis);
    }

    public void destroyObject(PooledObject<ShardedJedis> pooledShardedJedis) throws Exception {
        final ShardedJedis shardedJedis = pooledShardedJedis.getObject();
        for (Jedis jedis : shardedJedis.getAllShards()) {
            try {
                 jedis.quit();
                jedis.disconnect();
            } catch (Exception e) {
                log.error("ShardedJedisFactory destroyObject error",e);
            }
        }
    }

    public boolean validateObject(PooledObject<ShardedJedis> pooledShardedJedis) {
        try {
            ShardedJedis jedis = pooledShardedJedis.getObject();
            String pingResult;
            for (Jedis shard : jedis.getAllShards()) {
                pingResult = shard.ping();
                if (!"PONG".equals(pingResult)) {
                    log.error("testOnBorrow Ping Fail " + shard.getClient().getHost() + ":" + shard.getClient().getPort() + "pingResult : " + pingResult);
                    return false;
                }
            }
            return true;
        } catch (Exception ex) {
            log.error("testOnBorrow exception : ", ex);
            return false;
        }
    }

    public void activateObject(PooledObject<ShardedJedis> p) throws Exception {

    }

    public void passivateObject(PooledObject<ShardedJedis> p) throws Exception {

    }
}
